package testcase;

import org.opengauss.util.GT;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import util.DatetimeUtil;

import java.sql.*;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;

/**
 * @author zhoucong
 * @ClassName SavepointRollBackV4Demo
 * @Description TODO
 * @createTime 2023年06月15日15:24:00
 */
public class RollbackDemo {
    private static String URL = "jdbc:opengauss://113.137.58.68:50019/tpccdb?loggerLevel=OFF&socketTimeoutInConnecting=0&socketTimeout=0";
    private static final String USERNAME = "tpccuser";
    private static final String PASSWORD = "Huawei@123";

    private static final int IDLE_CONNECTIONS = 1000;

    private static final int COMMIT_CONNECTIONS = 500;

    private static final int ROLLBACK_CONNECTIONS = 500;

    private static int INSERT_NUM = 30;

    private static final Logger LOGGER = LoggerFactory.getLogger(RollbackDemo.class);

    private static void initTables() throws SQLException {
        Statement statement;
        // 初始化1000个链接
        try (Connection connection = DriverManager.getConnection(URL, USERNAME, PASSWORD)) {
            statement = connection.createStatement();
            statement.execute("drop schema if exists savepoint_test CASCADE");
            statement.execute("create schema savepoint_test");
            for (int i = 0; i < ROLLBACK_CONNECTIONS; i++) {
                statement.execute("create table if not exists savepoint_test.rollback_t" + i + "(id int, value varchar)");
            }
            for (int i = 0; i < COMMIT_CONNECTIONS; i++) {
                statement.execute("create table if not exists savepoint_test.commit_t" + i + "(id int, value varchar)");
            }
            for (int i = 0; i < IDLE_CONNECTIONS; i++) {
                statement.execute("create table if not exists savepoint_test.idle_t" + i + "(id int, value varchar)");
            }
            statement.close();
        }
        LOGGER.info(GT.tr("init tables complete, idle connections: {0}, commit connections: {1}, " +
                "rollback connections: {2}", IDLE_CONNECTIONS, COMMIT_CONNECTIONS, ROLLBACK_CONNECTIONS));
    }

    private static void executeInsertAndSavepoint(Connection connection, String tableName, int insertNum) throws SQLException {
        PreparedStatement preparedStatement = null;
        for (int i = 0; i < insertNum; i++) {
            try {
                connection.setSavepoint(GT.tr("savepoint{0}_{1}", tableName, String.valueOf(i)));
                preparedStatement = connection.prepareStatement(GT.tr("insert into savepoint_test.{0} values (?,?)", tableName));
                preparedStatement.setInt(1,i);
                preparedStatement.setString(2, "sfsfsfsfsfs");
                preparedStatement.execute();
            } finally {
                assert preparedStatement != null;
                preparedStatement.close();
            }
        }
    }

    // 创建1000长连接，都处于idle in transaction
    private static void executeIdleConnections () throws SQLException, InterruptedException {
        List<Connection> list = new ArrayList<>();
        try {
            int insertNum = 1000;
            for (int i = 0; i < IDLE_CONNECTIONS; i++) {
                Connection connection = DriverManager.getConnection(URL, USERNAME, PASSWORD);
                connection.setAutoCommit(false);
                executeInsertAndSavepoint(connection, GT.tr("idle_t{0}",i), INSERT_NUM);
                list.add(connection);
            }
            LOGGER.info(GT.tr("Create idle connections, connections: {0}.", IDLE_CONNECTIONS));
            // 死循环，连接长时间处于idle状态。
            while (true) {
                Thread.sleep(1000 * 60 * 60);
            }
        } finally {
            for(Connection connection : list) {
                connection.close();
            }
        }
    }

    // 创建1000长连接，反复提交事务
    private static void executeCommitConnections(int tableId) throws SQLException {
        Connection connection = DriverManager.getConnection(URL, USERNAME, PASSWORD);
        while (true) {
            try {
                connection.setAutoCommit(false);
                long startTime = System.currentTimeMillis();
                executeInsertAndSavepoint(connection, GT.tr("commit_t{0}", tableId), INSERT_NUM);
                connection.commit();
                long endTime = System.currentTimeMillis();
                LOGGER.info(GT.tr("Transaction commit，id: {0}, date: {1}, time interval: {2}", tableId, DatetimeUtil.getTheTimeInMilliseconds(), endTime - startTime));
            } catch (SQLException e) {
                throw new RuntimeException(e);
            }
        }
    }

    private static void executeRollbackConnections(int tableId) throws SQLException {
        Connection connection = DriverManager.getConnection(URL, USERNAME, PASSWORD);
        while (true) {
            try {
                connection.setAutoCommit(false);
                long startTime = System.currentTimeMillis();
                executeInsertAndSavepoint(connection, GT.tr("commit_t{0}", tableId), INSERT_NUM);
                connection.rollback();
                long endTime = System.currentTimeMillis();
                LOGGER.info(GT.tr("Transaction rollback，id: {0}, date: {1}, time interval: {2}", tableId, DatetimeUtil.getTheTimeInMilliseconds(), endTime - startTime));
            } catch (SQLException e) {
                throw new RuntimeException(e);
            }
        }
    }

    private static void execute() {
        ExecutorService executorService = Executors.newFixedThreadPool(COMMIT_CONNECTIONS + ROLLBACK_CONNECTIONS + 1);
        executorService.submit(()->{
            try {
                LOGGER.info(GT.tr("Start idle transaction thread"));
                executeIdleConnections();
            } catch (SQLException | InterruptedException e) {
                e.printStackTrace();
                throw new RuntimeException(e);
            }
        });

        for (int i = 0; i < COMMIT_CONNECTIONS; i++) {
            int finalI = i;
            executorService.submit(()->{
                LOGGER.info(GT.tr("Start commit thread {0}", finalI));
                try {
                    executeCommitConnections(finalI);
                } catch (SQLException  e) {
                    e.printStackTrace();
                    throw new RuntimeException(e);
                }
            });
        }

        for (int i = 0; i< ROLLBACK_CONNECTIONS; i++) {
            int finalI = i;
            executorService.submit(()->{
                LOGGER.info(GT.tr("Start rollback thread {0}", finalI));
                try {
                    executeRollbackConnections(finalI);
                } catch (SQLException  e) {
                    e.printStackTrace();
                    throw new RuntimeException(e);
                }
            });
        }
    }

    public static void main(String[] args) throws SQLException {
        if(args.length >=2 ) {
            URL = args[0];
            INSERT_NUM = Integer.parseInt(args[1]);
        }
        LOGGER.info(GT.tr("Start new thread {0}"));
        initTables();
        execute();
    }
}
